Skip to content

Handle explicit sharding in array truncation - #79

Closed
sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix-explicit-sharded-array-truncation
Closed

sylvesterkaczmarek wants to merge 1 commit into
google-deepmind:mainfrom
sylvesterkaczmarek:fix-explicit-sharded-array-truncation

Conversation

@sylvesterkaczmarek

@sylvesterkaczmarek sylvesterkaczmarek commented Aug 15, 2026 •

Copy link
Copy Markdown

Fixes #70.

Explicitly-sharded JAX arrays can reject ordinary slicing when the truncated dimension is not divisible by the mesh axis. Treescope currently hits that failure in _truncate_part_with_slices before the JIT out_shardings constraint can make the final result replicated.

For explicit NamedSharding, convert the first shrinking slice to a gather with an explicitly replicated output on the same mesh. Remaining slices then operate on a replicated intermediate, while non-explicit and older JAX paths keep the existing behavior.

The regression test forces four CPU devices in a subprocess and exercises JAXArrayAdapter.get_array_data_with_truncation with an explicit P("tp", None) sharding, verifying the 7x7 edge data and validity mask.

Local validation:

  • reproduced the original ShardingTypeError with four explicit CPU shards
  • patched truncation returned the expected 7x7 data and mask
  • ordinary non-explicit truncation returned the expected data and mask
  • Python syntax compilation passed

The full repository CI toolchain was not run locally because the environment could not fetch the development dependencies. The current Unittests workflow requires maintainer approval.

@sylvesterkaczmarek
sylvesterkaczmarek marked this pull request as ready for review August 15, 2026 18:33
@sylvesterkaczmarek

Copy link
Copy Markdown
Author

Following up after the maintenance pass. This branch is current with main, mergeable, has no failing checks or unresolved review threads, and the diff is unchanged from before today's signature-only rewrite. The new head is GitHub Verified. Could it be reviewed when convenient?

@sylvesterkaczmarek

Copy link
Copy Markdown
Author

Following up on #79, #80, #78. Could a maintainer take a look when available? Please let me know if any need changes or a different scope.

@danieldjohnson

Copy link
Copy Markdown
Collaborator

Closing this PR since the contributions appear to be automated and unrelated to real-world use of this library.

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

Array visualization fails for large, sharded arrays

2 participants